import read_mnist_data as data
import numpy as np
import sys

train_images = data.load_train_images()
X = train_images.reshape([60000, 784])

test_images = data.load_test_images()
test_labels = data.load_test_labels()

print("数据加载完成")


k = 10

def kmeans_train(X, k):
    n = X.shape[0]
    m = X.shape[1]
    cluster = []
    g_list = [[] for i in range(k)]

    for i in range(k):
        g_list[i].append(X[i])
        cluster.append(X[i])

    #是否是第一次运行，第一次要从k开始，之后从0开始
    first = True
    while True:
        if first:
            s = k
        else:
            s = 0
        for i in range(s, n):
            # 计算欧氏距离的平方
            # d记录最大的欧氏距离，l记录与当前样本向量 i距离最近的cluster分类的
            d = sys.maxsize
            l = 0
            for j in range(k):
                distance = np.sum(np.square(X[i] - cluster[j]))
                if distance < d:
                    d = distance
                    l = j
            ## X[i]应该分到 l 类，
            g_list[l].append(X[i])

        # 重新计算中心
        new_cluster = []
        for i in range(k):
            t = np.zeros(m)
            for sample in g_list[i]:
                t = t + sample
            t = t / len(g_list[i])
            new_cluster.append(t)

        #判断俩次族向量是否相同，如果相同停止分类
        stop = True
        for i in range(k):
            for j in range(m):
                if cluster[i][j] != new_cluster[i][j]:
                    stop = False
        if stop:
            break
        else:
            g_list = [[] for i in range(k)]
        cluster = new_cluster
        first = False
    return cluster

def test_kmeans(cluster):
    #cluster是训练结果表示k = 10类的各个类的中心点
    test_res = [[] for i in range(k)]
    test_data = test_images[:30].reshape([30, 784])


    for i in range(30):
        #计算距离
        d = sys.maxsize
        l = 0
        for j in range(k):
            distance = np.sum(np.square(test_data[i] - cluster[j]))
            if distance < d:
                d = distance
                l = j
        test_res[l].append(test_labels[i])
    return test_res
if __name__ == "__main__":
    cluster = kmeans_train(X, k)
    res = test_kmeans(cluster)
    print(res)